"""
# edge_visualizer.py

## 项目概述
本脚本用于从一系列 BMP 图像中提取右边缘坐标，并保存到 CSV 文件中，同时绘制折线图。该脚本主要用于图像处理和分析，特别是在需要跟踪图像中特定特征（如右边缘）的变化时。

## 架构描述
1. **ImageCropper 类**：
   - 用于裁剪图像。
   - 包含初始化方法 `__init__` 和三个主要方法：`initialize_with_reference` 和 `crop_image`。

2. **detect_vertical_lines 函数**：
   - 用于检测图像中的垂直线。
   - 使用 Sobel 算子和全局阈值进行边缘检测，并通过形态学操作增强垂直线条。

3. **visualize_edge_coordinates 函数**：
   - 主要函数，负责处理一系列 BMP 图像文件，提取右边缘坐标并保存到 CSV 文件中，同时绘制折线图。
   - 读取参考帧以初始化 `ImageCropper` 对象，并遍历所有图像文件进行处理。

## 使用方法
1. 确保输入路径中包含 BMP 格式的图像文件。
2. 运行脚本，提取右边缘坐标并生成 CSV 和折线图文件。

## 注意事项
- 该脚本假设输入的图像文件是 BMP 格式。
- 如果图像文件格式不同，需要修改 `visualize_edge_coordinates` 函数中的读取方式。
"""
import cv2
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm
import logging

logging.basicConfig(level=logging.ERROR, format='%(asctime)s - %(levelname)s - %(message)s')

class ImageCropper:
    def __init__(self):
        self.crop_width = None  # 裁剪宽度
        self.reference_right_edge = None  # 参考右边缘位置

    def initialize_with_reference(self, reference_img):
        # 转换为灰度图
        gray = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
        
        # 增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)

        # 检测垂直线
        binary = detect_vertical_lines(gray)

        # 找到右边缘
        right_edge = find_right_edge(binary)
        if right_edge is not None:
            self.reference_right_edge = right_edge
            # 设置裁剪宽度（从右边缘向左的距离）
            self.crop_width = right_edge - 100  # 可以根据需要调整这个值

    def crop_image(self, img):
        if self.crop_width is None:
            return img

        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)

        # 检测垂直线
        binary = detect_vertical_lines(gray)

        # 找到当前帧的右边缘
        current_right_edge = find_right_edge(binary)

        if current_right_edge is None:
            return img

        # 计算裁剪区域，使用动态边缘位置
        start_x = max(0, current_right_edge - self.crop_width)
        end_x = current_right_edge

        # 裁剪图像
        cropped = img[:, start_x:end_x]

        return cropped

def find_right_edge(binary_img, threshold_factor=2.0, local_region_offset_before=5, local_region_offset_after=5):
    # 计算垂直投影
    projection = np.sum(binary_img, axis=0)

    # 使用一维高斯卷积减少噪声影响
    window = 25 # 增大窗口大小
    sigma = window/3.0
    x = np.linspace(-window//2, window//2, window)
    gaussian = np.exp(-(x**2)/(2*sigma**2))
    gaussian = gaussian/np.sum(gaussian)

    # 应用高斯滤波平滑投影
    smoothed_projection = np.convolve(projection, gaussian, mode='same')

    # 计算梯度
    gradient = np.gradient(smoothed_projection)

    # 从右向左扫描找到第一个显著边缘
    threshold = np.std(gradient) * threshold_factor
    for i in range(len(gradient)-1, window, -1):
        if abs(gradient[i]) > threshold:
            # 使用局部最大值优化边缘位置
            local_region = gradient[max(0, i-local_region_offset_before):min(len(gradient), i+local_region_offset_after+1)]
            local_offset = np.argmax(np.abs(local_region))
            return i - local_region_offset_before + local_offset
    return None

def detect_vertical_lines(gray_img):
    # 使用Sobel算子增强垂直边缘检测
    sobelx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=5)
    sobelx = cv2.convertScaleAbs(sobelx)

    # 使用全局阈值
    _, binary = cv2.threshold(sobelx, 100, 255, cv2.THRESH_BINARY)

    # 形态学操作增强垂直线条
    kernel_vertical = np.ones((15,1), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel_vertical)
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel_vertical)

    return binary

def visualize_edge_coordinates(input_path, output_csv_path="edge_coordinates.csv", output_plot_path="edge_visualization.png"):
    image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
    if not image_files:
        logging.error(f"错误：在 {input_path} 中没有找到BMP文件")
        return
    image_files.sort()

    first_image_path = os.path.join(input_path, image_files[0])
    reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    if reference_frame is None:
        logging.error(f"错误：无法读取参考图片 {first_image_path}")
        return

    cropper = ImageCropper()
    cropper.initialize_with_reference(reference_frame)

    edge_coordinates = []
    for i, image_file in tqdm(enumerate(image_files), total=len(image_files), desc="处理图片"):
        img_path = os.path.join(input_path, image_file)
        current_frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        if current_frame is None:
            logging.error(f"错误：无法读取图片 {img_path}")
            edge_coordinates.append(None) # 记录错误帧
            continue

        gray = cv2.cvtColor(current_frame, cv2.COLOR_BGR2GRAY)
        binary = detect_vertical_lines(gray)
        right_edge = find_right_edge(binary)

        if right_edge is not None:
            relative_coordinate = right_edge - cropper.reference_right_edge if cropper.reference_right_edge is not None else 0
            edge_coordinates.append(relative_coordinate)
        else:
            logging.error(f"警告：未在图片 {image_file} 中检测到右边缘")
            edge_coordinates.append(None) # 记录未检测到边缘的情况

    # 转换为 Pandas DataFrame 并保存到 CSV
    df = pd.DataFrame({'Frame': range(len(image_files)), 'Relative_Edge_Coordinate': edge_coordinates})
    df.to_csv(output_csv_path, index=False)
    print(f"边缘坐标数据已保存到: {output_csv_path}")

    # 绘制折线图
    plt.figure(figsize=(12, 6))
    valid_indices = [i for i, coord in enumerate(edge_coordinates) if coord is not None]
    valid_coordinates = [coord for coord in edge_coordinates if coord is not None]
    plt.plot(valid_indices, valid_coordinates)
    plt.xlabel('Frame Number')
    plt.ylabel('Relative Right Edge Coordinate')
    plt.title('Relative Right Edge Coordinate vs. Frame Number')
    plt.grid(True)
    plt.savefig(output_plot_path)
    print(f"边缘坐标可视化图已保存到: {output_plot_path}")
    plt.close()


if __name__ == '__main__':
    input_path = r"F:\李双成\B-1\200-4"  # 替换为您的图像输入路径
    visualize_edge_coordinates(input_path)